'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 3_study_paralle_edge.py
* @Time: 2025/11/4
* @All Rights Reserve By Brtc
'''
from time import sleep
from typing import Any

from langchain_core.messages import AIMessage, HumanMessage
from langgraph.graph import StateGraph, MessagesState

graph_builder = StateGraph(MessagesState)

def chat_bot(state:MessagesState)->Any:
    return {"messages":[AIMessage(content="你好我是OpenAI 的机器人")]}

def paralle1(state:MessagesState)->Any:
    print("\n-----并行1\n")
    return {"messages":[HumanMessage(content="这是并行函数1")]}

def paralle2(state:MessagesState)->Any:
    print("\n-----并行2\n")
    sleep(3)
    return {"messages":[HumanMessage(content="这是并行函数2")]}


def chat_end(state:MessagesState)->Any:
    print("聊天结束")
    return {"messages":[HumanMessage(content="聊天结束函数")]}

#添加节点
graph_builder.add_node("chat_bot", chat_bot)
graph_builder.add_node("paralle1", paralle1)
graph_builder.add_node("paralle2", paralle2)
graph_builder.add_node("chat_end", chat_end)

#添加边
graph_builder.set_entry_point("chat_bot") # add stat
graph_builder.set_finish_point("chat_end")
graph_builder.add_edge("chat_bot", "paralle1")
graph_builder.add_edge("chat_bot", "paralle2")
graph_builder.add_edge("paralle1", "chat_end")

graph_app = graph_builder.compile()
print(graph_app.invoke({"messages":[HumanMessage(content="你好你是？")]}))